import argparse
import os

from basic.Dataloader import read_corpus


def main():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--model_path', default="bert-base-uncased")
    argparser.add_argument('--data', default="STAC_sp/train.json")
    argparser.add_argument('--max_edu_num', default="30")

    args, extra_args = argparser.parse_known_args()

    # tokenizer = AutoTokenizer.from_pretrained(plm_path)
    instances = read_corpus(args.data, max_edu_num=int(args.max_edu_num))

    out_file = args.data + ".raw"
    count = 0
    with open(out_file, mode='w') as outf:
        for instance in instances:
            for edu in instance.EDUs:
                count += 1
                text = edu['text']
                if text == "": text = "null"
                outf.write(text + "\n")

    print("write raw:", out_file) 
    print("line :", count) 

if __name__ == "__main__":
    main()